import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from PIL import Image, ImageOps
import matplotlib.pyplot as plt

def load_model(model_path='model.h5'):
    """加载训练好的模型"""
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型文件 {model_path} 不存在")
    
    # 禁用GPU以避免兼容性问题
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    
    try:
        model = tf.keras.models.load_model(model_path)
        print("模型加载成功")
        return model
    except Exception as e:
        print(f"模型加载失败: {e}")
        return None

def preprocess_image(image_path, target_size=(160, 160)):
    """预处理图像以匹配模型输入要求"""
    try:
        # 加载图像
        img = load_img(image_path, target_size=target_size)
        
        # 确保图像是RGB格式
        if img.mode != 'RGB':
            img = img.convert('RGB')
        
        # 转换为数组
        img_array = img_to_array(img)
        
        # 归一化
        img_array = img_array / 255.0
        
        # 添加批次维度
        img_array = np.expand_dims(img_array, axis=0)
        
        return img, img_array
    except Exception as e:
        print(f"图像加载失败: {e}")
        return None, None

def predict_and_visualize(model, original_img, img_array, display=True):
    """对图像进行分割预测并可视化结果"""
    try:
        # 执行预测
        prediction = model.predict(img_array)
        
        # 获取分割掩码
        mask = np.argmax(prediction[0], axis=-1)
        
        # 创建灰度分割结果
        gray_mask = create_grayscale_mask(mask)
        
        if display:
            # 显示原始图像和分割结果
            fig, axes = plt.subplots(1, 2, figsize=(12, 6))
            
            # 显示原始图像
            axes[0].imshow(original_img)
            axes[0].set_title('原始图像')
            axes[0].axis('off')
            
            # 显示灰度分割结果
            axes[1].imshow(gray_mask, cmap='gray')
            axes[1].set_title('分割结果(灰度)')
            axes[1].axis('off')
            
            plt.tight_layout()
            plt.show()
        
        return mask, gray_mask
    except Exception as e:
        print(f"预测过程出错: {e}")
        return None, None

def create_grayscale_mask(mask):
    """为分割掩码创建灰度表示"""
    # 定义灰度值映射
    gray_values = [0, 85, 170, 255]
    
    # 创建灰度掩码
    height, width = mask.shape
    gray_mask = np.zeros((height, width), dtype=np.uint8)
    
    # 应用灰度值到掩码
    for i in range(min(len(gray_values), np.max(mask) + 1)):
        gray_mask[mask == i] = gray_values[i]
    
    return gray_mask

def main():
    # 指定要分割的图像路径
    image_path = 'D:/project/dataset/segdata/images/Abyssinian_46.jpg'
    
    # 检查图像文件是否存在
    if not os.path.exists(image_path):
        print(f"图像文件 {image_path} 不存在")
        return
    
    # 加载模型
    model = load_model()
    if model is None:
        return
    
    # 预处理图像
    original_img, img_array = preprocess_image(image_path)
    if original_img is None or img_array is None:
        return
    
    # 执行预测并显示结果
    mask, gray_mask = predict_and_visualize(model, original_img, img_array)
    
    if mask is not None:
        # 保存灰度分割结果
        gray_img = Image.fromarray(gray_mask, mode='L')
        gray_img.save('segmentation_result_gray.png')
        print("灰度分割结果已保存为 segmentation_result_gray.png")

if __name__ == '__main__':
    main()